from PIL import Image
import numpy as np


def detect_color(image_path, lower_hsv, upper_hsv):
    # 加载图像并转换为RGB
    img = Image.open(image_path).convert('RGB')
    rgb_array = np.array(img)

    # 将RGB转换为HSV（手动实现）
    def rgb_to_hsv(rgb):
        r, g, b = rgb[..., 0], rgb[..., 1], rgb[..., 2]
        r, g, b = r / 255.0, g / 255.0, b / 255.0

        mx = np.max(rgb / 255.0, axis=-1)
        mn = np.min(rgb / 255.0, axis=-1)
        delta = mx - mn

        h = np.zeros_like(mx)
        cond = [delta == 0, (mx == r) & (delta != 0),
                (mx == g) & (delta != 0), (mx == b) & (delta != 0)]

        choices = [0,
                   60 * ((g - b) / delta) % 360,
                   60 * ((b - r) / delta + 2),
                   60 * ((r - g) / delta + 4)]
        h = np.select(cond, choices, default=0)
        s = np.where(mx == 0, 0, delta / mx) * 100
        v = mx * 100
        return np.stack([h, s, v], axis=-1)

    # 将整个图像转换为HSV
    hsv_image = np.apply_along_axis(rgb_to_hsv, 2, rgb_array)

    # 创建颜色掩码
    mask = np.ones(rgb_array.shape[:2], dtype=bool)
    for i in range(3):
        mask &= (hsv_image[:, :, i] >= lower_hsv[i])
        mask &= (hsv_image[:, :, i] <= upper_hsv[i])

    # 创建结果图像
    result = np.zeros_like(rgb_array)
    result[mask] = rgb_array[mask]

    return Image.fromarray(result), mask


# 定义黄色范围（H:30-60度, S:40-100%, V:40-100%）
lower_yellow = np.array([30, 40, 40])
upper_yellow = np.array([60, 100, 100])

# 执行检测并显示结果
result_img, mask = detect_color(r"C:\Users\songz\Desktop\Snipaste_2025-06-13_09-58-57.png", lower_yellow, upper_yellow)
result_img.show(title="黄色区域检测结果")
